import numpy as np
import pandas as pd
import csv
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D


# ============================================= 读取数据 ============================================
# 读取3d的.csv
def read_from_csv(path_csv):
    x_coords = []
    y_coords = []
    z_values = []

    with open(path_csv, 'r') as csvfile:
        reader = csv.DictReader(csvfile)
        for row in reader:
            x_coords.append(float(row['X']))
            y_coords.append(float(row['Y']))
            z_values.append(float(row['Z']))

    x_coords = np.unique(np.array(x_coords))
    y_coords = np.unique(np.array(y_coords))
    z_values = np.array(z_values).reshape((len(x_coords), len(y_coords)))

    return x_coords, y_coords, z_values

# ============================================= 画3D图（3D, 2D热力图, 2D等高线） ============================================
# 3d画图
def plot_3d_scatter(x_coords, y_coords, z_matrix):
    fig = plt.figure(figsize=(10, 8))
    ax = fig.add_subplot(111, projection='3d')

    X, Y = np.meshgrid(x_coords, y_coords)
    Z = z_matrix.T

    ax.scatter(X, Y, Z, c=Z, cmap='gist_rainbow')
    # ax.set_xlabel('X Coordinate')
    # ax.set_ylabel('Y Coordinate')
    # ax.set_zlabel('Z Value')
    # plt.title('3D Scatter Plot')
    # 隐藏轴标签和刻度
    # ax.set_xticks([])
    # ax.set_yticks([])
    # ax.set_zticks([])
    # 隐藏轴标签
    ax.set_xlabel('')
    ax.set_ylabel('')
    ax.set_zlabel('')
    # 显示网格线
    ax.grid(True)
    plt.show()


# 3d图的2d热力图（Heatmap）
def plot_2d_heatmap(x_coords, y_coords, z_matrix):
    fig, ax = plt.subplots(figsize=(10, 8))
    # 创建网格
    X, Y = np.meshgrid(x_coords, y_coords)
    Z = z_matrix.T
    # 绘制热力图
    cax = ax.imshow(Z, extent=[x_coords.min(), x_coords.max(), y_coords.min(), y_coords.max()], origin='lower',
                    cmap='gist_rainbow')
    # 添加颜色条
    cbar = fig.colorbar(cax)
    cbar.set_label('Z Value')
    ax.set_xlabel('X Coordinate')
    ax.set_ylabel('Y Coordinate')
    plt.title('2D Heatmap')
    plt.show()


# 3d图的2d等高线图（Contour Plot）
def plot_2d_contour(x_coords, y_coords, z_matrix):
    fig, ax = plt.subplots(figsize=(10, 8))
    # 创建网格
    X, Y = np.meshgrid(x_coords, y_coords)
    Z = z_matrix.T
    # 绘制等高线图
    contour = ax.contour(X, Y, Z, levels=20, cmap='gist_rainbow')
    # 添加颜色条
    plt.clabel(contour, inline=True, fontsize=8)
    ax.set_xlabel('X Coordinate')
    ax.set_ylabel('Y Coordinate')
    plt.title('2D Contour Plot')
    plt.show()





# ============================================= 单波束画图 ============================================
# 画比较图 -- 3d的eh折线
def plot_compare_3d_eh(x, line1, line2, label1, lable2, xlabel, ylabel):
    # 创建图形和轴
    plt.figure(figsize=(14, 6))
    # 绘制第一个数据集的折线图
    plt.plot(x, line1, label=label1, color='red', linewidth=2)
    # 绘制第二个数据集的折线图
    plt.plot(x, line2, label=lable2, color='blue', linewidth=2, linestyle='--')
    # 设置标签和标题
    plt.xlabel(xlabel)
    plt.ylabel(ylabel)
    # plt.title('XY Line Plot of Two Datasets')
    # 添加图例
    plt.legend()
    # 显示图形
    plt.grid(True)
    plt.show()


# 画3d的eh折线
def plot_linechart_3d_eh(x, line, label, xlabel, ylabel):
    # 创建图形和轴
    plt.figure(figsize=(14, 6))
    # 绘制第一个数据集的折线图
    plt.plot(x, line, label=label, color='red', linewidth=2)
    # 设置标签和标题
    plt.xlabel(xlabel)
    plt.ylabel(ylabel)
    # plt.title('XY Line Plot of Two Datasets')
    # 添加图例
    plt.legend()
    # 显示图形
    plt.grid(True)
    plt.show()


# 找到最大值所在行列
def format_z2eh(z_matrix):
    z_matrix = np.exp(z_matrix / 20)
    z_matrix = z_matrix / np.max(z_matrix)
    z_matrix = 20 * np.log10(z_matrix)
    indices = (z_matrix < -20).nonzero()
    z_matrix[indices] = -20
    # 找到 z_matrix 最大值元素的位置
    max_pos = np.unravel_index(np.argmax(z_matrix), z_matrix.shape)
    row_idx, col_idx = max_pos
    # 取 z_matrix 所在行和列的数据
    z_matrix_max_row = z_matrix[row_idx, :]
    z_matrix_max_col = z_matrix[:, col_idx]
    #
    print("z_matrix 最大值位置:", max_pos)
    # print("z_matrix 最大值所在行数据:\n", z_matrix_max_row)
    # print("z_matrix 最大值所在列数据:\n", z_matrix_max_col)
    return max_pos, z_matrix_max_row, z_matrix_max_col


def compare_beam1_3d_e_h():
    path_csv_1 = "../files/testdataset/1/1-3D.csv"
    path_csv_2 = "../files/testdataset/2/2-3D.csv"
    # 读取CSV文件
    x_coords_1, y_coords_1, z_matrix_1 = read_from_csv(path_csv_1)
    x_coords_2, y_coords_2, z_matrix_2 = read_from_csv(path_csv_2)
    #
    max_pos_1, z_matrix_max_row_1, z_matrix_max_col_1 = format_z2eh(z_matrix_1)
    max_pos_2, z_matrix_max_row_2, z_matrix_max_col_2 = format_z2eh(z_matrix_2)
    #
    plot_compare_3d_eh(x_coords_1, z_matrix_max_row_1, z_matrix_max_row_2, "AGA", "BP", "degree", "normalized pattern")
    plot_compare_3d_eh(y_coords_1, z_matrix_max_col_1, z_matrix_max_col_2, "AGA", "BP", "degree", "normalized pattern")


def linechart_beam1_3d_e_h():
    path_csv = "../files/testdataset/1/1-3D.csv"
    # 读取CSV文件
    x_coords, y_coords, z_matrix = read_from_csv(path_csv)
    #
    max_pos, z_matrix_max_row, z_matrix_max_col = format_z2eh(z_matrix)
    #
    plot_linechart_3d_eh(x_coords, z_matrix_max_row, "AGA", "degree", "normalized pattern")
    plot_linechart_3d_eh(y_coords, z_matrix_max_col, "AGA", "degree", "normalized pattern")



#===================================================== 双波束画图 =================================================
# 找到前n个位置
def find_top_n_positions(matrix, n=2):
    # 获取矩阵展平后的索引，并按降序排序
    flat_indices = np.argsort(matrix, axis=None)[::-1]
    # 获取前n个最大值的位置
    positions = [np.unravel_index(idx, matrix.shape) for idx in flat_indices[:n]]
    return positions


def process_matrix(matrix, pos1, pos2):
    list1 = []  # 存 0~pos1
    list2 = []  # 存 pos1~pos2
    list3 = []  # 存 pos2~135
    # list2
    x = pos1[0]
    y = pos1[1]
    while x <= pos2[0] and y <= pos2[1]:
        if x == pos1[0]:
            list2.append(matrix[pos1])
        elif x == pos2[0]:
            list2.append(matrix[pos2])
        else:
            list2.append(matrix[x, y])
        x = x + 1
        y = y + 1
    # list1
    x = pos1[0] - 1
    y = pos1[1] - 1
    while 0 <= x and 0 <= y:
        list1.append(matrix[x, y])
        x = x - 1
        y = y - 1
    # 反转 list1
    list1.reverse()
    # list3
    x = pos2[0] + 1
    y = pos2[1] + 1
    while x < 135 and y < 135:
        list3.append(matrix[x, y])
        x = x + 1
        y = y + 1
    line = list1 + list2 + list3
    return np.array(line)


# 归一化z_matrix
def format_z_matrix(z_matrix):
    z_matrix = np.exp(z_matrix / 20)
    z_matrix = z_matrix / np.max(z_matrix)
    z_matrix = 20 * np.log10(z_matrix)
    indices = (z_matrix < -20).nonzero()
    z_matrix[indices] = -20
    return z_matrix


def compare_beam2_3d_e_h():
    path_csv_1 = "../files/testdataset/10/10-3D.csv"
    path_csv_2 = "../files/testdataset/11/11-3D.csv"
    # 读取CSV文件
    x_coords_1, y_coords_1, z_matrix_1 = read_from_csv(path_csv_1)
    x_coords_2, y_coords_2, z_matrix_2 = read_from_csv(path_csv_2)
    #
    z_matrix_1 = format_z_matrix(z_matrix_1)
    z_matrix_2 = format_z_matrix(z_matrix_2)
    # 找到 z_matrix_1 最大值排前2的元素位置
    positions_1 = find_top_n_positions(z_matrix_1, n=2)
    pos1_1, pos2_1 = positions_1
    # 找到 z_matrix_2 最大值排前2的元素位置
    positions_2 = find_top_n_positions(z_matrix_2, n=2)
    pos1_2, pos2_2 = positions_2
    # 将 pos1 和 pos2 连线上的点保存到数组 line
    line_1 = process_matrix(z_matrix_1, pos1_2, pos1_1)
    line_2 = process_matrix(z_matrix_2, pos1_2, pos1_1)
    # 绘制 line 的折线图
    plot_compare_3d_eh(x_coords_1[2:-2], line_1, line_2, "AGA", "VS", "degree", "normalized pattern")
    plot_compare_3d_eh(y_coords_1[2:-2], line_1, line_2, "AGA", "VS", "degree", "normalized pattern")


def linechart_beam2_3d_e_h():
    path_csv = "../files/testdataset/10/10-3D.csv"
    # 读取CSV文件
    x_coords, y_coords, z_matrix = read_from_csv(path_csv)
    #
    z_matrix = format_z_matrix(z_matrix)
    # 找到 z_matrix 最大值排前2的元素位置
    positions = find_top_n_positions(z_matrix, n=2)
    pos1, pos2 = positions
    # 将 pos1 和 pos2 连线上的点保存到数组 line
    line = process_matrix(z_matrix, pos1, pos1)
    # 绘制 line 的折线图
    x, label, xlabel, ylabel = x_coords[1:-1], "AGA", "degree", "normalized pattern"
    # 创建图形和轴
    plt.figure(figsize=(14, 6))
    # 绘制第一个数据集的折线图
    plt.plot(x, line, label=label, color='red', linewidth=2)
    # 设置标签和标题
    plt.xlabel(xlabel)
    plt.ylabel(ylabel)
    # plt.title('XY Line Plot of Two Datasets')
    # 添加图例
    plt.legend()
    # 显示图形
    plt.grid(True)
    plt.show()




# 示例调用函数
if __name__ == "__main__":
    # 读取并画图3d
    x_coords, y_coords, z_matrix = read_from_csv("../files/testdataset/2/2-3D.csv")
    z_matrix = np.exp(z_matrix / 20)
    z_matrix = z_matrix / np.max(z_matrix)
    z_matrix = 20 * np.log10(z_matrix)
    indices = (z_matrix < -20).nonzero()
    z_matrix[indices] = -20
    #
    plot_3d_scatter(x_coords, y_coords, z_matrix)
    # plot_2d_heatmap(x_coords, y_coords, z_matrix)
    # plot_2d_contour(x_coords, y_coords, z_matrix)
    #
    # 比较3d的横纵折线
    # compare_beam2_3d_e_h()
    # compare_beam1_3d_e_h()
    #
    # 单个3d的eh折线图
    # linechart_beam1_3d_e_h()
    # linechart_beam2_3d_e_h()


